torchrun  --nproc_per_node=4 mnist_project.py --epochs=3
